import os
from pinecone import Pinecone, ServerlessSpec
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from tqdm import tqdm
import logging
import time

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

def setup_pinecone():
    """初始化Pinecone连接"""
    api_key = 'pcsk_3Hx5py_TFFCuUMU2bVyo3V9yQy4tzpMuotrB4irA2brsi3nncJRkN1Ut7CHUjFm1x8oXn9'
    
    try:
        pc = Pinecone(api_key=api_key)
        logging.info("Pinecone初始化成功")
        return pc
    except Exception as e:
        logging.error(f"Pinecone初始化失败: {e}")
        raise

def create_mnist_index(pc):
    """创建MNIST索引并上传数据"""
    index_name = "mnist-handwritten-digits"
    
    # 如果索引已存在，先删除
    existing_indexes = pc.list_indexes()
    if index_name in [index.name for index in existing_indexes.indexes]:
        logging.info(f"删除已存在的索引: {index_name}")
        pc.delete_index(index_name)
        time.sleep(10)  # 等待索引完全删除
    
    # 创建新索引
    logging.info(f"创建新索引: {index_name}")
    pc.create_index(
        name=index_name,
        dimension=64,  # MNIST数据是8x8=64维
        metric="cosine",
        spec=ServerlessSpec(
            cloud='aws',
            region='us-east-1'
        )
    )
    time.sleep(10)  # 等待索引准备就绪
    
    return pc.Index(index_name)

def upload_training_data(index, X_train, y_train):
    """上传训练数据到Pinecone"""
    logging.info("开始上传训练数据...")
    
    vectors = []
    batch_size = 100
    
    for i in tqdm(range(len(X_train)), desc="上传数据"):
        vector_id = f"digit_{i}"
        vector = X_train[i].astype(float).tolist()
        metadata = {"label": int(y_train[i])}
        vectors.append({
            "id": vector_id,
            "values": vector,
            "metadata": metadata
        })
        
        # 分批上传
        if len(vectors) >= batch_size:
            index.upsert(vectors=vectors)
            vectors = []
    
    # 上传剩余数据
    if vectors:
        index.upsert(vectors=vectors)
    
    logging.info(f"成功创建索引，并上传了 {len(X_train)} 条数据")

def test_accuracy(index, X_test, y_test, k=11):
    """测试Pinecone索引的准确率"""
    logging.info(f"开始测试准确率 (k={k})...")
    
    predictions = []
    
    for i in tqdm(range(len(X_test)), desc="测试准确率"):
        query_vector = X_test[i].astype(float).tolist()
        
        # 查询Pinecone
        results = index.query(
            vector=query_vector,
            top_k=k,
            include_metadata=True
        )
        
        # 统计最近邻的标签
        neighbor_labels = [match['metadata']['label'] for match in results['matches']]
        
        # 使用多数投票决定预测标签
        predicted_label = max(set(neighbor_labels), key=neighbor_labels.count)
        predictions.append(predicted_label)
    
    accuracy = accuracy_score(y_test, predictions)
    logging.info(f"当k={k}时，使用Pinecone的准确率: {accuracy:.4f} ({accuracy*100:.2f}%)")
    
    return accuracy

def main():
    """主函数"""
    logging.info("开始MNIST Pinecone训练流程")
    
    # 加载数据
    logging.info("加载MNIST数据集")
    digits = load_digits()
    X, y = digits.data, digits.target
    
    # 划分训练测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    
    logging.info(f"训练集大小: {len(X_train)}, 测试集大小: {len(X_test)}")
    
    try:
        # 初始化Pinecone
        pc = setup_pinecone()
        
        # 创建索引
        index = create_mnist_index(pc)
        
        # 上传训练数据
        upload_training_data(index, X_train, y_train)
        
        # 测试准确率
        accuracy = test_accuracy(index, X_test, y_test, k=11)
        
        logging.info("训练流程完成")
        
        # 可选：清理索引
        # pc.delete_index("mnist-handwritten-digits")
        
    except Exception as e:
        logging.error(f"流程执行失败: {e}")

if __name__ == "__main__":
    main()